#include "ekf.h"
#include <time.h>
#include <stdio.h>

EkfSlam::EkfSlam(Mat (*scanner)(int), Mat qi, Mat si)
{
	q = qi.clone();
	s = si.clone();
	scan = scanner;
	mState = (Mat_<double>(3,1) << 0, 0, 0);
	Prr = (Mat_<double>(3,3) << 0,0,0,0,0,0,0,0,0);

	time_t timer;
	time(&timer);

	r = new RNG(timer);
}

EkfSlam::~EkfSlam()
{
	delete r;
}

// Transition Function
Mat EkfSlam::fv(Mat R, Mat u, Mat n)
{
	double p=u.at<double>(0);
	double dth=u.at<double>(1);

	double x=R.at<double>(0);
	double y=R.at<double>(1);
	double theta=R.at<double>(2);

	p += n.at<double>(0);
	double ntheta = theta + dth + n.at<double>(1);
	if(ntheta > M_PI)
		ntheta -= 2*M_PI;
	if(ntheta < -M_PI)
		ntheta += 2*M_PI;

	return (Mat_<double>(3,1) <<
		p*cos(theta)+x,
		p*sin(theta)+y,
		ntheta
	);
}

Mat EkfSlam::df_dxb(Mat R, Mat u, Mat n)
{
	double p=u.at<double>(0);
	double dth=u.at<double>(1);

	double x=R.at<double>(0);
	double y=R.at<double>(1);
	double theta=R.at<double>(2);

	double n1 = n.at<double>(0);
	double n2 = n.at<double>(1);

	return (Mat_<double>(3,3) <<
		1, 0, -(p+n1)*sin(theta),
		0, 1, (p+n1)*cos(theta),
		0, 0, 1
	);
}

Mat EkfSlam::df_dn(Mat R, Mat u, Mat n)
{
	double p=u.at<double>(0);
	double dth=u.at<double>(1);

	double x=R.at<double>(0);
	double y=R.at<double>(1);
	double theta=R.at<double>(2);

	double n1 = n.at<double>(0);
	double n2 = n.at<double>(1);

	return (Mat_<double>(3,2) <<
		-sin(theta), 0,
		cos(theta),0,
		0,1
	);
}

void EkfSlam::move(Mat R, Mat u, Mat n)
{
		Mat Rtemp = fv(R, u, n);

		// get jacobians
		Mat dfdr = df_dxb(R, u,n);
		Mat dfdn = df_dn(R, u,n);

		mState.at<double>(0) = Rtemp.at<double>(0);
		mState.at<double>(1) = Rtemp.at<double>(1);
		mState.at<double>(2) = Rtemp.at<double>(2);

		Mat q2 = q.clone();
		q2.at<double>(0) *= q2.at<double>(0);
		q2.at<double>(1) *= q2.at<double>(1);
		Mat N = Mat::diag(q2);

		// update covariance matrix

	if(Prm.rows) {
		Prr = dfdr*Prr*dfdr.t() + dfdn*N*dfdn.t();
		Prm = dfdr*Prm;
		Pmr = Prm.t();
	}
}

// Observation Function
Mat EkfSlam::h(Mat R, Mat Li)
{
	double x = R.at<double>(0);	// robot x
	double y = R.at<double>(1);	// robot y
	double theta = R.at<double>(2);	//robot theta

	Mat RT = (Mat_<double>(2,2) << cos(theta), sin(theta),
				       -sin(theta), cos(theta)
		);

	Mat TT = (Mat_<double>(2,1) << x,y);

	TT = RT*(Li-TT);

	double px = TT.at<double>(0);
	double py = TT.at<double>(1);

	return (Mat_<double>(2,1) <<
		sqrt(pow(px,2)+pow(py,2)),
		atan2(py,px)
	);
}

Mat EkfSlam::dh_dxb(Mat R, Mat Li)
{
	double x = R.at<double>(0);	// robot x
	double y = R.at<double>(1);	// robot y
	double theta = R.at<double>(2);	//robot theta

	double xw = Li.at<double>(0); // Landmark x
	double yw = Li.at<double>(1); // Landmark y

	double pr = pow(xw-x,2) + pow(yw-y,2);

	return (Mat_<double>(2,3) <<
		(x-xw)/sqrt(pr), (y-yw)/sqrt(pr), 0,
		(yw-y)/pr, (x-xw)/pr, -1
	);
}

Mat EkfSlam::dh_dLi(Mat R, Mat Li)
{
	double x = R.at<double>(0); // robot x
	double y = R.at<double>(1); // robot y

	double xw = Li.at<double>(0); // Landmark x
	double yw = Li.at<double>(1); // Landmark y

	double pr = pow(xw-x,2)+pow(yw-y,2);

	return (Mat_<double>(2,2) <<
		(xw-x)/sqrt(pr), (yw-y)/sqrt(pr),
		(y-yw)/pr, (xw-x)/pr
	);
}

void EkfSlam::correct(Mat yi, Mat v, int i)
{
	if(!yi.rows)
		return;

	double x = mState.at<double>(0); // robot x
	double y = mState.at<double>(1); // robot y
	double theta = mState.at<double>(2); //robot theta

	Mat Li = (Mat_<double>(2,1) <<
		mState.at<double>(2*i+3),
		mState.at<double>(2*i+4)
	);

	//////////////////////////////////////////////////////////////
	// setup handy data for kalman correction

	Mat Hr = dh_dxb(mState(Rect(0,0,1,3)), Li);
	Mat Hl = dh_dLi(mState(Rect(0,0,1,3)), Li);

	Mat Hr_t = Hr.t();
	Mat Hl_t = Hl.t();

	Mat H;
	H.push_back(Hr_t);
	H.push_back(Hl_t);
	H = H.t();

	Mat s2 = s.clone();
	s2.at<double>(0) *= s2.at<double>(0);
	s2.at<double>(1) *= s2.at<double>(1);

	Mat R = Mat::diag(s2);

	/////////////////////////////////////////////////////////////
	// Kalman Correction
	Mat zbar = yi-h(mState(Rect(0,0,1,3)), Li);

	if(zbar.at<double>(1) > M_PI)
		zbar.at<double>(1) -= 2*M_PI;
	if(zbar.at<double>(1) < -M_PI)
		zbar.at<double>(1) += 2*M_PI;
	///////////////////////////////////////
	Mat Prl = Prm(Rect(2*i,0,2,3)).clone();
	Mat Plr = Pmr(Rect(0,2*i,3,2)).clone();
	Mat Pll = Pmm(Rect(2*i,2*i,2,2)).clone();

	Mat Pl;
	Mat pTemp1, pTemp2;
	Pl = Prr.t();
	pTemp1 = Prl.t();
	Pl.push_back(pTemp1);
	Pl = Pl.t();

	pTemp1 = Plr.t();
	pTemp2 = Pll.t();
	pTemp1.push_back(pTemp2);
	pTemp1 = pTemp1.t();
	Pl.push_back(pTemp1);
	///////////////////////////////////////
	Mat Z = H*Pl*H.t()+R;

	Mat t = zbar.t()*Z.inv()*zbar;

	// update map cov
	// In large maps with small uncertainties you may want to limit
	// the updates to within a certain range of uncertainty with an
	// if statement around this next block of code.
//	if(t.at<double>(0) < 9) {
		Mat P = Prr.t();
		P.push_back(Plr);
		P = P.t();
		Mat t2 = Prm.clone();
		Mat Plm = Pmm(Rect(0,2*i,Pmm.cols,2)).clone();
		t2.push_back(Plm);
		t2 = t2.t();
		P.push_back(t2);

		Mat K = P*H.t()*Z.inv();
		mState += K*zbar;
		Mat Pmod = K*Z*K.t();
		Prr -= Pmod(Rect(0,0,3,3));
		Prm -= Pmod(Rect(3,0,Pmod.cols-3,3));
		//Pmr = Prm.t();
		Pmr -= Pmod(Rect(0,3,3,Pmod.rows-3));
		Pmm -= Pmod(Rect(3,3,Pmod.cols-3,Pmod.rows-3));
//	}
}

// inverse observation model
Mat EkfSlam::g(Mat R, Mat yi)
{
	// rotate and translate measurement into world space

	double pyi = yi.at<double>(0);		// range
	double thetay = yi.at<double>(1);	// bearing

	double x = R.at<double>(0);
	double y = R.at<double>(1);
	double theta = R.at<double>(2);

	Mat Lr = (Mat_<double>(3,1) <<		// homogenous coords
		pyi*cos(thetay),		// x,y local
		pyi*sin(thetay),
		1
	);

	Mat Rt = (Mat_<double>(2,3) <<
		cos(theta), -sin(theta), x,
		sin(theta), cos(theta), y
	);

	return Rt*Lr;
}

Mat EkfSlam::dg_dxb(Mat R, Mat yi)
{
	double p = yi.at<double>(0);		// observation range
	double thetay = yi.at<double>(1);	// observation bearing

	double px = p*cos(thetay);
	double py = p*sin(thetay);

	double theta = R.at<double>(2);

	return (Mat_<double>(2,3) <<
		1, 0, -px*sin(theta) - py*cos(theta),
		0, 1, px*cos(theta)  - py*sin(theta)
	);
}

Mat EkfSlam::dg_dyi(Mat R, Mat yi)
{
	double p = yi.at<double>(0);		// observation range
	double thetay = yi.at<double>(1);	// observation bearing

	double px = p*cos(thetay);
	double py = p*sin(thetay);

	double theta = R.at<double>(2);

	return (Mat_<double>(2,2) <<
		cos(thetay)*cos(theta)-sin(theta)*sin(thetay), -py*cos(theta)-px*sin(theta),
		sin(theta)*cos(thetay)+cos(theta)*sin(thetay), -py*sin(theta)+px*cos(theta)
	);
}


void EkfSlam::addFeature(Mat yi, int i)
{
	printf("Discovered Landmark %d\n", i);
	Mat s2 = s.clone();
	s2.at<double>(0) *= s2.at<double>(0);
	s2.at<double>(1) *= s2.at<double>(1);

	Mat R = Mat::diag(s2);

	Mat Gr = dg_dxb(mState(Rect(0,0,1,3)), yi);
	Mat Gy = dg_dyi(mState(Rect(0,0,1,3)), yi);

	Mat L = g(mState(Rect(0,0,1,3)), yi);
	mState.push_back(L);

	Mat Pll = Gr*Prr*Gr.t() + Gy*R*Gy.t();
	Mat Plr = Gr*Prr;
	Mat Plm;

	if(!i) {
		Pmr = Plr.clone();
		Prm = Pmr.t();
		Pmm = Pll.clone();
	}
	else {
		Plm = Gr*Prm;
		Pmr.push_back(Plr);
		Prm = Pmr.t();

		// append Pml to right of Pmm
		Pmm = Pmm.t();
		Pmm.push_back(Plm);
		Pmm = Pmm.t();
		// append Pll to right of Plm
		Plm = Plm.t();
		Pll = Pll.t();
		Plm.push_back(Pll);
		Plm = Plm.t();

		// append [Plm|Pll] to bottom of Pmm
		Pmm.push_back(Plm);
	}
}

bool EkfSlam::getMap(Mat& map, Mat &cov)
{
	if(mState.rows > 3) {
		map = mState(Rect(0, 3, 1, mState.rows-3)).clone();
		cov = Pmm.clone();
		return true;
	}
	return false;
}

Mat EkfSlam::getRobot()
{
	return mState(Rect(0,0,1,3)).clone();
}

Mat EkfSlam::step(Mat u)
{
	Mat observations;
	Mat n = (Mat_<double>(2,1) << 0, 0);
	Mat v = (Mat_<double>(2,1) << 0, 0);
	Mat Y;

	// motion prediction updates covariance and state vector
	move(mState(Rect(0,0,1,3)), u, n);

	// landmark correction
	for(int i=0; i < (mState.rows-3)/2; i++) {
		v.at<double>(0) = s.at<double>(0)*r->gaussian(.5);
		v.at<double>(1) = s.at<double>(1)*r->gaussian(.5);

		Mat yi = scan(i);
		if(yi.rows) {
			observations.push_back(yi);
			yi+=v;
			correct(scan(i), v, i);
		}
	}

	// initialize any new landmarks discovered
	for(int i=(mState.rows-3)/2;;i++) { 
		Mat f = scan(i);
		if(f.rows == 2) {
			addFeature(f,i);
			observations.push_back(f);
		}
		else
			break;
	}

	// return raw sensor data for simple navigation tasks
	return observations;
}
